#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from std_msgs.msg import Float64MultiArray
from sensor_msgs.msg import JointState
from functools import partial
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

class MotorStateBridge(Node):
    def __init__(self):
        qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=10
        )
        super().__init__('motor_state_bridge')
        self.declare_parameter('ugv_ids', ['ugv1'])
        self.ugv_ids = self.get_parameter('ugv_ids').value
        self.pub = {}
        for ugv_id in self.ugv_ids:
            sub_topic = f'/world/default/model/{ugv_id}/joint_state'
            pub_topic = f'/{ugv_id}/actuators/chassis/out/motor_data'
            self.sub = self.create_subscription(
            JointState,
            sub_topic,
            partial(self.joint_state_callback, ugv_id=ugv_id),
            qos
            )
        

            self.pub[ugv_id] = self.create_publisher(
                Float64MultiArray,
                pub_topic,
                10
            )


 
    def joint_state_callback(self, msg: JointState, ugv_id: str):
        name_list = msg.name
        velocity_list = msg.velocity
        effort_list = msg.effort

        try:
            i_rl = name_list.index("rover_ackermann/RearLeftWheelJoint")
            i_rr = name_list.index("rover_ackermann/RearRightWheelJoint")

            speed_rl = velocity_list[i_rl]
            speed_rr = velocity_list[i_rr]
            torque_rl = effort_list[i_rl]
            torque_rr = effort_list[i_rr]

            out_msg = Float64MultiArray()
            out_msg.data = [speed_rl, speed_rr, torque_rl, torque_rr]

            self.pub[ugv_id].publish(out_msg)
            

        except ValueError:
            self.get_logger().warn("Rear joints not found in joint_state")

    

def main(args=None):
    rclpy.init(args=args)
    node = MotorStateBridge()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
